// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cantor/function/resolver.h>
#include <cantor/type/promote.h>
#include <cantor/function/signature.h>

namespace cantor {
    static void resolve_function_fast(const FunctionSet *set, uint64_t fid,
                                      cantor::pb::catalog::FunctionResolveResult &result, std::vector<int> &variadics) {
        for (auto i = 0; i < set->functions.size(); i++) {
            auto &it = set->functions[i];
            if (it->is_variadic()) {
                variadics.push_back(i);
            }
            if (it->base_id() == fid) {
                result.set_found(true);
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_NONE) {
                    result.set_scalar_id(it->id());
                }
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_SIMD) {
                    result.set_simd_id(it->id());
                }
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_GPU) {
                    result.set_gpu_id(it->id());
                }
            }
        }
    }

    static void resolve_function_variadic(const FunctionSet *set, const cantor::pb::catalog::FunctionResolveArgs &args,
                                          const std::vector<int> &variadics,
                                          cantor::pb::catalog::FunctionResolveResult &result) {
        cantor::pb::catalog::FunctionDefine def;
        def.set_func_name(args.func_name());
        auto type_size = args.args().size();
        def.mutable_args()->Reserve(type_size);
        for (auto i = 0; i < type_size - 1; i++) {
            *def.mutable_args()->Add() = args.args()[i];
        }
        def.set_is_variadic(true);
        *def.mutable_variadic_type() = args.args(type_size - 1);
        auto rs = make_function_signature_no_feature(&def);
        result.set_found(false);
        if (!rs.ok()) {
            return;
        }
        for (auto i = 0; i < variadics.size(); i++) {
            auto &it = set->functions[variadics[i]];
            if (def.base_id() == it->base_id()) {
                result.set_found(true);
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_NONE) {
                    result.set_scalar_id(it->id());
                }
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_SIMD) {
                    result.set_simd_id(it->id());
                }
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_GPU) {
                    result.set_gpu_id(it->id());
                }
            }
        }
    }

    static bool resolve_function_define_promote(const cantor::pb::catalog::FunctionDefine *def,
                                         const cantor::pb::catalog::FunctionResolveArgs &args,
                                         cantor::pb::catalog::FunctionResolveResult &result) {
        int total_cost = 0;
        for (auto i = 0; i < def->args_size(); i++) {
            auto &from = args.args()[i];
            auto &to = def->args()[i];
            if (type_equal(from, to)) {
                continue;
            }
            auto cost = promote_cost(from, to);
            if (cost > 0) {
                total_cost += cost;
                continue;
            }
            return false;
        }
        result.set_found(true);
        switch (def->feature()) {
            case cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_NONE:
                result.set_scalar_id(def->id());
                result.set_scalar_score(total_cost);
                break;
            case cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_SIMD:
                result.set_simd_id(def->id());
                result.set_simd_score(total_cost);
                break;
            case cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_GPU:
                result.set_gpu_id(def->id());
                result.set_gpu_score(total_cost);
                break;
        }

        return true;
    }

    static bool resolve_function_define_promote_variadic(const cantor::pb::catalog::FunctionDefine *def,
                                         const cantor::pb::catalog::FunctionResolveArgs &args,
                                         cantor::pb::catalog::FunctionResolveResult &result) {
        int total_cost = 0;
        for (auto i = 0; i < def->args_size(); i++) {
            auto &from = args.args()[i];
            auto &to = def->args()[i];
            if (type_equal(from, to)) {
                continue;
            }
            auto cost = promote_cost(from, to);
            if (cost > 0) {
                total_cost += cost;
                continue;
            }
            return false;
        }

        auto &variadic_type = def->variadic_type();
        for (auto i = def->args_size(); i <args.args_size();  i++) {
            auto &from = args.args()[i];
            if (type_equal(from, variadic_type)) {
                continue;
            }
            auto cost = promote_cost(from, variadic_type);
            if (cost > 0) {
                total_cost += cost;
                continue;
            }
            return false;

        }

        result.set_found(true);
        switch (def->feature()) {
            case cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_NONE:
                result.set_scalar_id(def->id());
                result.set_scalar_score(total_cost);
                break;
            case cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_SIMD:
                result.set_simd_id(def->id());
                result.set_simd_score(total_cost);
                break;
            case cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_GPU:
                result.set_gpu_id(def->id());
                result.set_gpu_score(total_cost);
                break;
        }

        return true;
    }


    static void resolve_function_promote(const FunctionSet *set, const cantor::pb::catalog::FunctionResolveArgs &args,
                                         cantor::pb::catalog::FunctionResolveResult &result) {
        std::vector<cantor::pb::catalog::FunctionResolveResult> results;
        for (auto i = 0; i < set->functions.size(); i++) {
            if (set->functions[i]->args().size() == args.args().size()) {
                cantor::pb::catalog::FunctionResolveResult r;
                r.set_gpu_score(std::numeric_limits<int64_t>::max());
                r.set_scalar_score(std::numeric_limits<int64_t>::max());
                r.set_simd_score(std::numeric_limits<int64_t>::max());
                if (resolve_function_define_promote(set->functions[i].get(), args, r)) {
                    results.push_back(r);
                }
                continue;
            }

            if ((set->functions[i]->args().size() < args.args().size()) && set->functions[i]->is_variadic() && set->functions[i]->has_variadic_type() ) {
                cantor::pb::catalog::FunctionResolveResult r;
                r.set_gpu_score(std::numeric_limits<int64_t>::max());
                r.set_scalar_score(std::numeric_limits<int64_t>::max());
                r.set_simd_score(std::numeric_limits<int64_t>::max());
                if (resolve_function_define_promote_variadic(set->functions[i].get(), args, r)) {
                    results.push_back(r);
                }
                continue;
            }

        }

        if (results.empty()) {
            result.set_found(false);
            return;
        }
        result.set_found(true);
        result.set_gpu_score(std::numeric_limits<int64_t>::max());
        result.set_scalar_score(std::numeric_limits<int64_t>::max());
        result.set_simd_score(std::numeric_limits<int64_t>::max());
        for (auto i = 0; i < results.size(); i++) {
            auto &r = results[i];
            if (r.has_scalar_id()) {
                if (r.scalar_score() < result.scalar_score()) {
                    result.set_scalar_score(r.scalar_score());
                    result.set_scalar_id(r.scalar_id());
                }
            }
            if (r.has_simd_id()) {
                if (r.simd_score() < result.simd_score()) {
                    result.set_simd_score(r.simd_score());
                    result.set_simd_id(r.simd_id());
                }
            }
            if (r.has_gpu_id()) {
                if (r.gpu_score() < result.gpu_score()) {
                    result.set_gpu_score(r.gpu_score());
                    result.set_gpu_id(r.gpu_id());
                }
            }
        }
        result.set_found(true);
    }

    void resolve_function(const FunctionRegistry *registry, const cantor::pb::catalog::FunctionResolveArgs &args,
                          cantor::pb::catalog::FunctionResolveResult &result) {
        auto f_set = registry->find_function_set(args.func_name());
        if (!f_set) {
            result.set_found(false);
            return;
        }
        auto fid = make_function_args_signature(args);
        std::vector<int> variadics;
        resolve_function_fast(f_set.get(), fid, result, variadics);
        if (result.found()) {
            return;
        }
        /// resolve slow
        if (!variadics.empty()) {
            resolve_function_variadic(f_set.get(), args, variadics, result);
        }
        if (result.found()) {
            return;
        }
        /// promote
        resolve_function_promote(f_set.get(), args, result);
    }
} // namespace cantor
